package ai.mspbots.trace.config;

import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import brave.Tracer;

//@Component
public class TraceIdGatewayFilterFactory extends AbstractGatewayFilterFactory<Object> {

    private final Tracer tracer;

    public TraceIdGatewayFilterFactory(Tracer tracer) {
        this.tracer = tracer;
    }

    @Override
    public GatewayFilter apply(Object config) {
        return (exchange, chain) -> {
            String traceId = tracer.currentSpan() != null ? tracer.currentSpan().context().traceIdString() : null;
            if (traceId != null) {
                exchange.getRequest().mutate()
                    .header("X-B3-TraceId", traceId)
                    .build();
            }
            return chain.filter(exchange);
        };
    }
}